import logging
import socket
import os
import json
from data_tranform_server import data_tcp_port, data_udp_port

char_set = 'ABCDEFGHIGKLMNOPQRSTUVWXYZabcdefghigklmnopqrstuvwxyz0123456789'
char_set_len = 62


def tcp_recevier(host_ip):
    tcp_server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp_server_socket.bind((host_ip, 5000))
    while True:
        try:
            tcp_server_socket.listen(64)
            client_socket = tcp_server_socket.accept()
            recv_data = client_socket[0].recv(8192)
            req = json.loads(recv_data.decode('utf-8'))
            if req['method'] == 'traceroute':
                resp = traceroute(req['args'][0])
            elif req['method'] == 'send_tcp_file':
                resp = send_tcp_file(req['args'][0], req['args'][1])
            elif req['method'] == 'send_udp_file':
                resp = send_udp_file(req['args'][0], req['args'][1])
            resp_bytes = json.dumps(resp)
            client_socket[0].send(resp_bytes.encode('utf-8'))
        except Exception as e:
            logging.error("TCP Error:" + str(e))


def traceroute(dst_ip: str) -> dict:
    try:
        result_reader = os.popen("traceroute %s" % dst_ip)
        result = result_reader.read()
        splt = result.split('\n')
        if len(splt) > 1:
            ans = []
            for item in splt[1:-1]:
                ans.append(item.split())
            return {
                'code': 0,
                'message': 'unreachable',
                'data': ans
            }
        return {
            'code': -2,
            'message': 'unreachable',
            'data': None
        }
    except Exception as e:
        return {
            'code': -1,
            'message': str(e),
            'data': None
        }


def gen_random_bytes(size: int) -> bytes:
    loop = size // char_set_len
    left = size % char_set_len
    res_str = ""
    for i in range(loop):
        res_str += char_set
    if char_set_len > 0:
        res_str += char_set[0:left]
    return res_str.encode()


def send_udp_file(host_ip: str, size_kb: int):
    udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    dest_addr = (host_ip, data_udp_port)
    send_data = gen_random_bytes(1024 * size_kb)
    udp_socket.sendto(send_data, dest_addr)


def send_tcp_file(host_ip: str, size_kb: int):
    tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp_socket.connect((host_ip, data_tcp_port))
    send_data = gen_random_bytes(1024 * size_kb)
    tcp_socket.send(send_data)
    tcp_socket.close()


if __name__ == '__main__':
    target_ip = input("Target IP: ")
    send_tcp_file(target_ip, 2048)
